

#include <stdint.h>
#include <string.h>

#include "list.h"


/* nothing */
#define DEBUG_ASSERT(x)
#define ROUNDUP(a, b) (((a) + ((b)-1)) & ~((b)-1))
#define ROUNDDOWN(a, b) ((a) & ~((b)-1))


struct free_heap_chunk {
    struct list_node node;
    size_t len;
};

struct heap {
    void *base;
    size_t len;
    size_t remaining;
    size_t low_watermark;
    struct list_node free_list;
};

// heap static vars
static struct heap theheap;

// structure placed at the beginning every allocation
struct alloc_struct_begin {
    void *ptr;
    size_t size;
};


// try to insert this free chunk into the free list, consuming the chunk by merging it with
// nearby ones if possible. Returns base of whatever chunk it became in the list.
static struct free_heap_chunk *heap_insert_free_chunk(struct free_heap_chunk *chunk) {
    uintptr_t chunk_end = (uintptr_t)chunk + chunk->len;

    struct free_heap_chunk *next_chunk;
    struct free_heap_chunk *last_chunk;

    theheap.remaining += chunk->len;

    // walk through the list, finding the node to insert before
    list_for_every_entry(&theheap.free_list, next_chunk, struct free_heap_chunk, node) {
        if (chunk < next_chunk) {
            DEBUG_ASSERT(chunk_end <= (uintptr_t)next_chunk);

            list_add_before(&next_chunk->node, &chunk->node);

            goto try_merge;
        }
    }

    // walked off the end of the list, add it at the tail
    list_add_tail(&theheap.free_list, &chunk->node);

    // try to merge with the previous chunk
try_merge:
    last_chunk = list_prev_type(&theheap.free_list, &chunk->node, struct free_heap_chunk, node);
    if (last_chunk) {
        if ((uintptr_t)last_chunk + last_chunk->len == (uintptr_t)chunk) {
            // easy, just extend the previous chunk
            last_chunk->len += chunk->len;

            // remove ourself from the list
            list_delete(&chunk->node);

            // set the chunk pointer to the newly extended chunk, in case
            // it needs to merge with the next chunk below
            chunk = last_chunk;
        }
    }

    // try to merge with the next chunk
    if (next_chunk) {
        if ((uintptr_t)chunk + chunk->len == (uintptr_t)next_chunk) {
            // extend our chunk
            chunk->len += next_chunk->len;

            // remove them from the list
            list_delete(&next_chunk->node);
        }
    }

    return chunk;
}


static struct free_heap_chunk *heap_create_free_chunk(void *ptr, size_t len, bool allow_debug)
{
    DEBUG_ASSERT((len % sizeof(void *)) == 0); // size must be aligned on pointer boundary
    DEBUG_ASSERT(len > sizeof(struct free_heap_chunk));

    struct free_heap_chunk *chunk = (struct free_heap_chunk *)ptr;
    chunk->len = len;

    return chunk;
}

void *miniheap_alloc( size_t size, unsigned int alignment )
{
    void *ptr;

    // alignment must be power of 2
    if ( alignment & (alignment - 1) )
        return NULL;

    // we always put a size field + base pointer + magic in front of the allocation
    size += sizeof(struct alloc_struct_begin);

    // make sure we allocate at least the size of a struct free_heap_chunk so that
    // when we free it, we can create a struct free_heap_chunk struct and stick it
    // in the spot
    if (size < sizeof(struct free_heap_chunk))
        size = sizeof(struct free_heap_chunk);

    // round up size to a multiple of native pointer size
    size = ROUNDUP(size, sizeof(void *));

    // deal with nonzero alignments
    if (alignment > 0) {
        if (alignment < 16)
            alignment = 16;

        // add alignment for worst case fit
        size += alignment;
    }

    // walk through the list
    ptr = NULL;
    struct free_heap_chunk *chunk;
    list_for_every_entry(&theheap.free_list, chunk, struct free_heap_chunk, node) {
        DEBUG_ASSERT((chunk->len % sizeof(void *)) == 0); // len should always be a multiple of pointer size

        // is it big enough to service our allocation?
        if (chunk->len >= size) {
            ptr = chunk;

            // remove it from the list
            struct list_node *next_node = list_next(&theheap.free_list, &chunk->node);
            list_delete(&chunk->node);

            if (chunk->len > size + sizeof(struct free_heap_chunk)) {
                // there's enough space in this chunk to create a new one after the allocation
                struct free_heap_chunk *newchunk = heap_create_free_chunk((uint8_t *)ptr + size, chunk->len - size, true);

                // truncate this chunk
                chunk->len -= chunk->len - size;

                // add the new one where chunk used to be
                if (next_node)
                    list_add_before(next_node, &newchunk->node);
                else
                    list_add_tail(&theheap.free_list, &newchunk->node);
            }

            // the allocated size is actually the length of this chunk, not the size requested
            DEBUG_ASSERT(chunk->len >= size);
            size = chunk->len;

            ptr = (void *)((uintptr_t)ptr + sizeof(struct alloc_struct_begin));

            // align the output if requested
            if (alignment > 0) {
                ptr = (void *)ROUNDUP((uintptr_t)ptr, (uintptr_t)alignment);
            }

            struct alloc_struct_begin *as = (struct alloc_struct_begin *)ptr;
            as--;
            as->ptr = (void *)chunk;
            as->size = size;
            theheap.remaining -= size;

            if (theheap.remaining < theheap.low_watermark) {
                theheap.low_watermark = theheap.remaining;
            }

            break;
        }
    }

    /**/
    return ptr;
}

void miniheap_free( void *ptr )
{
    if ( ptr == NULL )
        return;

    // check for the old allocation structure
    struct alloc_struct_begin *as = (struct alloc_struct_begin *)ptr;
    as--;

    // looks good, create a free chunk and add it to the pool
    heap_insert_free_chunk(heap_create_free_chunk(as->ptr, as->size, true));
}


void *miniheap_realloc(void *ptr, size_t size) {
    /* slow implementation */
    if ( ptr == NULL )
        return miniheap_alloc(size, 0);

    if (size == 0) {
        miniheap_free(ptr);
        return NULL;
    }

    // XXX better implementation
    void *p = miniheap_alloc(size, 0);
    if ( p == NULL )
        return NULL;

    memcpy(p, ptr, size); // XXX wrong
    miniheap_free(ptr);

    return p;
}



void miniheap_init(void *ptr, size_t len)
{

    // initialize the free list
    list_initialize(&theheap.free_list);

    // align the buffer to a ptr
    if (((uintptr_t)ptr % sizeof(uintptr_t)) > 0) {
        uintptr_t aligned_ptr = ROUNDUP((uintptr_t)ptr, sizeof(uintptr_t));

        DEBUG_ASSERT((aligned_ptr - (uintptr_t)ptr) < len);

        len -= aligned_ptr - (uintptr_t)ptr;
        ptr = (void *)aligned_ptr;
    }

    // set the heap range
    theheap.base = ptr;
    theheap.len = len;
    theheap.remaining = 0; // will get set by heap_insert_free_chunk()
    theheap.low_watermark = 0;

    // if passed a default range, use it
    if (len > 0)
        heap_insert_free_chunk(heap_create_free_chunk(ptr, len, true));
}


/**/
void *malloc( size_t size )
{
    return miniheap_alloc(size, 0);
}

